package com.gateway.plugin.tenant;

import com.gateway.common.dto.RuleData;
import com.gateway.common.dto.SelectorData;
import com.gateway.common.dto.convert.rule.impl.TenantRuleHandle;
import com.gateway.common.enums.PluginEnum;
import com.gateway.common.exception.ThrowingFunction;
import com.gateway.plugin.api.GatewayPluginChain;
import com.gateway.plugin.api.result.GatewayResultEnum;
import com.gateway.plugin.api.result.GatewayResultWrap;
import com.gateway.plugin.api.util.WebFluxResultUtils;
import com.gateway.plugin.base.AbstractGatewayPlugin;
import com.gateway.plugin.base.util.CacheKeyUtils;
import com.gateway.plugin.tenant.handle.TenantPluginDataHandler;
import com.gateway.plugin.tenant.strategy.JwtConvertStrategy;
import com.gateway.plugin.tenant.strategy.JwtConvertStrategyFactory;
import com.google.common.collect.ImmutableMap;
import io.jsonwebtoken.Jwt;
import io.jsonwebtoken.JwtParser;
import io.jsonwebtoken.JwtParserBuilder;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.SignatureAlgorithm;
import io.jsonwebtoken.security.Keys;
import org.apache.commons.lang3.StringUtils;
import org.springframework.util.CollectionUtils;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.Collection;
import java.util.Date;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * @program: api-gateway
 * @description:
 * @author: YuKai Fan
 * @create: 2025/2/14 09:55
 **/
public class TenantGatewayPlugin extends AbstractGatewayPlugin {

    private static final String TENANT = "X-Tenant";
    private static final String APP = "X-App";
    private static final String SIGNATURE = "X-Signature";

    @Override
    protected Mono<Void> doExecute(ServerWebExchange exchange, GatewayPluginChain chain, SelectorData selector, RuleData rule) {
        String tenantId = exchange.getRequest().getHeaders().getFirst(TENANT);
        String appId = exchange.getRequest().getHeaders().getFirst(APP);
        String signature = exchange.getRequest().getHeaders().getFirst(SIGNATURE);
        if (StringUtils.isAnyBlank(tenantId, appId, signature)) {
            Object error = GatewayResultWrap.error(exchange, GatewayResultEnum.MISS_PARAM);
            return WebFluxResultUtils.result(exchange, error);
        }
        if (Objects.isNull(rule)) {
            Object error = GatewayResultWrap.error(exchange, GatewayResultEnum.RULE_NOT_FOUND);
            return WebFluxResultUtils.result(exchange, error);
        }
        if (StringUtils.isBlank(rule.getHandle())) {
            Object error = GatewayResultWrap.error(exchange, GatewayResultEnum.RULE_NOT_FOUND);
            return WebFluxResultUtils.result(exchange, error);
        }
        TenantRuleHandle tenantRuleHandle = TenantPluginDataHandler.CACHE_HANDLE.get().obtainHandle(CacheKeyUtils.INST.getKey(rule));
        if (Boolean.TRUE.equals(tenantRuleHandle.getTenantDisable())) {
            Object error = GatewayResultWrap.error(exchange, GatewayResultEnum.TENANT_DISABLE);
            return WebFluxResultUtils.result(exchange, error);
        }
        if (Boolean.TRUE.equals(tenantRuleHandle.getAppDisable())) {
            Object error = GatewayResultWrap.error(exchange, GatewayResultEnum.APP_DISABLE);
            return WebFluxResultUtils.result(exchange, error);
        }
        Map<String, Object> jwtBody = checkAuthorization(signature, tenantRuleHandle.getAppSecret());
        if (CollectionUtils.isEmpty(jwtBody)) {
            Object error = GatewayResultWrap.error(exchange, GatewayResultEnum.SIGN_IS_NOT_PASS);
            return WebFluxResultUtils.result(exchange, error);
        }
        // check
        if (!checkBody(jwtBody, tenantId, appId)) {
            Object error = GatewayResultWrap.error(exchange, GatewayResultEnum.BAD_REQUEST_FOR_CLIENT);
            return WebFluxResultUtils.result(exchange, error);
        }
        return chain.execute(executeRuleHandle(tenantRuleHandle, exchange, jwtBody));
    }

    @Override
    public int getOrder() {
        return PluginEnum.TENANT.getCode();
    }

    @Override
    public String named() {
        return PluginEnum.TENANT.getName();
    }

    private boolean checkBody(Map<String, Object> jwtBody, String... params) {
        Collection<String> values = jwtBody.values().stream().map(String::valueOf).collect(Collectors.toSet());
        for (String param : params) {
            if (!values.contains(param)) {
                return false;
            }
        }
        return true;
    }

    private ServerWebExchange executeRuleHandle(final TenantRuleHandle tenantRuleHandle, final ServerWebExchange exchange, final Map<String, Object> jwtBody) {
        JwtConvertStrategy jwtConvertStrategy = JwtConvertStrategyFactory.newInstance(tenantRuleHandle.getHandleType());
        return jwtConvertStrategy.convert(tenantRuleHandle, exchange, jwtBody);
    }

    /**
     * check Authorization.
     *
     * @param authorization the authorization after processing
     * @param secretKey     secretKey of authorization
     * @return Map
     */
    private Map<String, Object> checkAuthorization(final String authorization,
                                                   final String secretKey) {

        if (StringUtils.isEmpty(authorization)) {
            return null;
        }
        JwtParserBuilder jwtParserBuilder = Jwts.parserBuilder();
        JwtParser jwtParser = jwtParserBuilder.build();
        if (jwtParser.isSigned(authorization)) {
            jwtParserBuilder.setSigningKey(secretKey.getBytes(StandardCharsets.UTF_8));
            JwtParser jwtParserExec = jwtParserBuilder
                    .build();
            Jwt jwt = ThrowingFunction.wrap(() -> jwtParserExec.parseClaimsJws(authorization));
            if (jwt == null) {
                return null;
            }
            return (Map<String, Object>) jwt.getBody();
        }
        return null;
    }

    public static void main(String[] args) {

        Map<String, Object> map = ImmutableMap.<String, Object>builder().put("tenantId", 123)
                .put("appId", 123)
                .build();
        String token = Jwts.builder()
                .setClaims(map)
                .setIssuedAt(new Date())
                .setExpiration(new Date(System.currentTimeMillis() + 30 * 1000))
                .signWith(Keys.hmacShaKeyFor("gateway-test-gateway-test-gateway-test".getBytes(StandardCharsets.UTF_8)), SignatureAlgorithm.HS256)
                .compact();
        System.out.println(token);
    }
}
